use async_trait::async_trait;
use sqlx::{PgPool, Row};
use crate::common::storage::*;
use chrono::{DateTime, Utc};
use uuid::Uuid;

pub struct PostgreSQLStorage {
    pool: PgPool,
}

impl PostgreSQLStorage {
    pub async fn new(database_url: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let pool = PgPool::connect(database_url).await?;
        
        // 初始化数据库表
        Self::init_tables(&pool).await?;
        
        Ok(Self { pool })
    }
    
    async fn init_tables(pool: &PgPool) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        // 创建节点表
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS nodes (
                id VARCHAR(255) PRIMARY KEY,
                hostname VARCHAR(255) NOT NULL,
                ip_address VARCHAR(45) NOT NULL,
                tags TEXT[] NOT NULL DEFAULT '{}',
                status VARCHAR(50) NOT NULL DEFAULT 'offline',
                last_heartbeat TIMESTAMP WITH TIME ZONE,
                capabilities TEXT[] NOT NULL DEFAULT '{}',
                metadata JSONB NOT NULL DEFAULT '{}',
                group_ids TEXT[] NOT NULL DEFAULT '{}',
                created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
                updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
            )
            "#
        ).execute(pool).await?;
        
        // 创建分组表
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS groups (
                id VARCHAR(255) PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                description TEXT NOT NULL DEFAULT '',
                selector JSONB NOT NULL,
                nodes TEXT[] NOT NULL DEFAULT '{}',
                created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
                updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
            )
            "#
        ).execute(pool).await?;
        
        // 创建命令记录表
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS commands (
                id VARCHAR(255) PRIMARY KEY,
                command_type VARCHAR(50) NOT NULL,
                target TEXT NOT NULL,
                payload JSONB NOT NULL,
                status VARCHAR(50) NOT NULL DEFAULT 'pending',
                created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
                completed_at TIMESTAMP WITH TIME ZONE,
                results JSONB NOT NULL DEFAULT '[]'
            )
            "#
        ).execute(pool).await?;
        
        // 创建文件元数据表
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS files (
                id VARCHAR(255) PRIMARY KEY,
                filename VARCHAR(255) NOT NULL,
                file_size BIGINT NOT NULL DEFAULT 0,
                checksum VARCHAR(255) NOT NULL DEFAULT '',
                status VARCHAR(50) NOT NULL DEFAULT 'pending',
                created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
                updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
            )
            "#
        ).execute(pool).await?;
        
        Ok(())
    }
}

#[async_trait]
impl Storage for PostgreSQLStorage {
    // 节点管理实现
    async fn register_node(&self, node: &NodeInfo) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            INSERT INTO nodes (id, hostname, ip_address, tags, status, last_heartbeat, capabilities, metadata, group_ids)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
            ON CONFLICT (id) DO UPDATE SET
                hostname = EXCLUDED.hostname,
                ip_address = EXCLUDED.ip_address,
                tags = EXCLUDED.tags,
                status = EXCLUDED.status,
                last_heartbeat = EXCLUDED.last_heartbeat,
                capabilities = EXCLUDED.capabilities,
                metadata = EXCLUDED.metadata,
                group_ids = EXCLUDED.group_ids,
                updated_at = NOW()
            "#
        )
        .bind(&node.id)
        .bind(&node.hostname)
        .bind(&node.ip_address)
        .bind(&node.tags)
        .bind(match node.status {
            NodeStatus::Online => "online",
            NodeStatus::Offline => "offline",
            NodeStatus::Maintenance => "maintenance",
            NodeStatus::Error => "error",
        })
        .bind(node.last_heartbeat)
        .bind(&node.capabilities)
        .bind(serde_json::to_value(&node.metadata)?)
        .bind(&node.group_ids)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn unregister_node(&self, node_id: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query("DELETE FROM nodes WHERE id = $1")
            .bind(node_id)
            .execute(&self.pool)
            .await?;
        
        Ok(())
    }
    
    async fn get_node(&self, node_id: &str) -> Result<Option<NodeInfo>, Box<dyn std::error::Error + Send + Sync>> {
        let row = sqlx::query(
            r#"
            SELECT id, hostname, ip_address, tags, status, last_heartbeat, capabilities, metadata, group_ids
            FROM nodes WHERE id = $1
            "#
        )
        .bind(node_id)
        .fetch_optional(&self.pool)
        .await?;
        
        if let Some(row) = row {
            let node = NodeInfo {
                id: row.get("id"),
                hostname: row.get("hostname"),
                ip_address: row.get("ip_address"),
                tags: row.get("tags"),
                status: match row.get::<&str, _>("status") {
                    "online" => NodeStatus::Online,
                    "offline" => NodeStatus::Offline,
                    "maintenance" => NodeStatus::Maintenance,
                    "error" => NodeStatus::Error,
                    _ => NodeStatus::Offline,
                },
                last_heartbeat: row.get("last_heartbeat"),
                capabilities: row.get("capabilities"),
                metadata: serde_json::from_value(row.get("metadata"))?,
                group_ids: row.get("group_ids"),
            };
            Ok(Some(node))
        } else {
            Ok(None)
        }
    }
    
    async fn get_all_nodes(&self) -> Result<Vec<NodeInfo>, Box<dyn std::error::Error + Send + Sync>> {
        let rows = sqlx::query(
            r#"
            SELECT id, hostname, ip_address, tags, status, last_heartbeat, capabilities, metadata, group_ids
            FROM nodes ORDER BY created_at DESC
            "#
        )
        .fetch_all(&self.pool)
        .await?;
        
        let mut nodes = Vec::new();
        for row in rows {
            let node = NodeInfo {
                id: row.get("id"),
                hostname: row.get("hostname"),
                ip_address: row.get("ip_address"),
                tags: row.get("tags"),
                status: match row.get::<&str, _>("status") {
                    "online" => NodeStatus::Online,
                    "offline" => NodeStatus::Offline,
                    "maintenance" => NodeStatus::Maintenance,
                    "error" => NodeStatus::Error,
                    _ => NodeStatus::Offline,
                },
                last_heartbeat: row.get("last_heartbeat"),
                capabilities: row.get("capabilities"),
                metadata: serde_json::from_value(row.get("metadata"))?,
                group_ids: row.get("group_ids"),
            };
            nodes.push(node);
        }
        
        Ok(nodes)
    }
    
    async fn update_node_status(&self, node_id: &str, status: NodeStatus) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let status_str = match status {
            NodeStatus::Online => "online",
            NodeStatus::Offline => "offline",
            NodeStatus::Maintenance => "maintenance",
            NodeStatus::Error => "error",
        };
        
        sqlx::query(
            "UPDATE nodes SET status = $1, updated_at = NOW() WHERE id = $2"
        )
        .bind(status_str)
        .bind(node_id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn update_node_heartbeat(&self, node_id: &str, timestamp: DateTime<Utc>) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            "UPDATE nodes SET last_heartbeat = $1, updated_at = NOW() WHERE id = $2"
        )
        .bind(timestamp)
        .bind(node_id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    // 分组管理实现（简化版本）
    async fn create_group(&self, group: &Group) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            INSERT INTO groups (id, name, description, selector, nodes)
            VALUES ($1, $2, $3, $4, $5)
            "#
        )
        .bind(&group.id)
        .bind(&group.name)
        .bind(&group.description)
        .bind(serde_json::to_value(&group.selector)?)
        .bind(&group.nodes)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn get_group(&self, group_id: &str) -> Result<Option<Group>, Box<dyn std::error::Error + Send + Sync>> {
        let row = sqlx::query(
            "SELECT id, name, description, selector, nodes FROM groups WHERE id = $1"
        )
        .bind(group_id)
        .fetch_optional(&self.pool)
        .await?;
        
        if let Some(row) = row {
            let group = Group {
                id: row.get("id"),
                name: row.get("name"),
                description: row.get("description"),
                selector: serde_json::from_value(row.get("selector"))?,
                nodes: row.get("nodes"),
                created_at: Utc::now(), // 简化处理
                updated_at: Utc::now(),
            };
            Ok(Some(group))
        } else {
            Ok(None)
        }
    }
    
    async fn get_all_groups(&self) -> Result<Vec<Group>, Box<dyn std::error::Error + Send + Sync>> {
        let rows = sqlx::query(
            "SELECT id, name, description, selector, nodes FROM groups ORDER BY created_at DESC"
        )
        .fetch_all(&self.pool)
        .await?;
        
        let mut groups = Vec::new();
        for row in rows {
            let group = Group {
                id: row.get("id"),
                name: row.get("name"),
                description: row.get("description"),
                selector: serde_json::from_value(row.get("selector"))?,
                nodes: row.get("nodes"),
                created_at: Utc::now(), // 简化处理
                updated_at: Utc::now(),
            };
            groups.push(group);
        }
        
        Ok(groups)
    }
    
    async fn update_group(&self, group: &Group) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            UPDATE groups SET name = $1, description = $2, selector = $3, nodes = $4, updated_at = NOW()
            WHERE id = $5
            "#
        )
        .bind(&group.name)
        .bind(&group.description)
        .bind(serde_json::to_value(&group.selector)?)
        .bind(&group.nodes)
        .bind(&group.id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn delete_group(&self, group_id: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query("DELETE FROM groups WHERE id = $1")
            .bind(group_id)
            .execute(&self.pool)
            .await?;
        
        Ok(())
    }
    
    async fn add_node_to_group(&self, node_id: &str, group_id: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            UPDATE groups SET nodes = array_append(nodes, $1) WHERE id = $2
            "#
        )
        .bind(node_id)
        .bind(group_id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn remove_node_from_group(&self, node_id: &str, group_id: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            UPDATE groups SET nodes = array_remove(nodes, $1) WHERE id = $2
            "#
        )
        .bind(node_id)
        .bind(group_id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn get_group_nodes(&self, group_id: &str) -> Result<Vec<NodeInfo>, Box<dyn std::error::Error + Send + Sync>> {
        let rows = sqlx::query(
            r#"
            SELECT id, hostname, ip_address, tags, status, last_heartbeat, capabilities, metadata, group_ids
            FROM nodes WHERE id = ANY(SELECT nodes FROM groups WHERE id = $1)
            "#
        )
        .bind(group_id)
        .fetch_all(&self.pool)
        .await?;
        
        let mut nodes = Vec::new();
        for row in rows {
            let node = NodeInfo {
                id: row.get("id"),
                hostname: row.get("hostname"),
                ip_address: row.get("ip_address"),
                tags: row.get("tags"),
                status: match row.get::<&str, _>("status") {
                    "online" => NodeStatus::Online,
                    "offline" => NodeStatus::Offline,
                    "maintenance" => NodeStatus::Maintenance,
                    "error" => NodeStatus::Error,
                    _ => NodeStatus::Offline,
                },
                last_heartbeat: row.get("last_heartbeat"),
                capabilities: row.get("capabilities"),
                metadata: serde_json::from_value(row.get("metadata"))?,
                group_ids: row.get("group_ids"),
            };
            nodes.push(node);
        }
        
        Ok(nodes)
    }
    
    // 命令管理实现（简化版本）
    async fn store_command(&self, command: &CommandRecord) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            INSERT INTO commands (id, command_type, target, payload, status, results)
            VALUES ($1, $2, $3, $4, $5, $6)
            "#
        )
        .bind(&command.id)
        .bind(match command.command_type {
            CommandType::Shell => "shell",
            CommandType::Script => "script",
            CommandType::ServiceRestart => "service_restart",
            CommandType::FileOperation => "file_operation",
        })
        .bind(&command.target)
        .bind(&command.payload)
        .bind(match command.status {
            CommandStatus::Pending => "pending",
            CommandStatus::Running => "running",
            CommandStatus::Completed => "completed",
            CommandStatus::Failed => "failed",
            CommandStatus::Cancelled => "cancelled",
        })
        .bind(serde_json::to_value(&command.results)?)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn update_command_status(&self, command_id: &str, status: CommandStatus) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let status_str = match status {
            CommandStatus::Pending => "pending",
            CommandStatus::Running => "running",
            CommandStatus::Completed => "completed",
            CommandStatus::Failed => "failed",
            CommandStatus::Cancelled => "cancelled",
        };
        
        sqlx::query(
            "UPDATE commands SET status = $1 WHERE id = $2"
        )
        .bind(status_str)
        .bind(command_id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn add_command_result(&self, command_id: &str, result: &CommandResult) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            UPDATE commands SET results = array_append(results, $1) WHERE id = $2
            "#
        )
        .bind(serde_json::to_value(result)?)
        .bind(command_id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn get_command(&self, command_id: &str) -> Result<Option<CommandRecord>, Box<dyn std::error::Error + Send + Sync>> {
        let row = sqlx::query(
            "SELECT id, command_type, target, payload, status, results FROM commands WHERE id = $1"
        )
        .bind(command_id)
        .fetch_optional(&self.pool)
        .await?;
        
        if let Some(row) = row {
            let command = CommandRecord {
                id: row.get("id"),
                command_type: match row.get::<&str, _>("command_type") {
                    "shell" => CommandType::Shell,
                    "script" => CommandType::Script,
                    "service_restart" => CommandType::ServiceRestart,
                    "file_operation" => CommandType::FileOperation,
                    _ => CommandType::Shell,
                },
                target: row.get("target"),
                payload: row.get("payload"),
                status: match row.get::<&str, _>("status") {
                    "pending" => CommandStatus::Pending,
                    "running" => CommandStatus::Running,
                    "completed" => CommandStatus::Completed,
                    "failed" => CommandStatus::Failed,
                    "cancelled" => CommandStatus::Cancelled,
                    _ => CommandStatus::Pending,
                },
                created_at: Utc::now(), // 简化处理
                completed_at: None,
                results: serde_json::from_value(row.get("results"))?,
            };
            Ok(Some(command))
        } else {
            Ok(None)
        }
    }
    
    async fn get_commands_by_target(&self, target: &str) -> Result<Vec<CommandRecord>, Box<dyn std::error::Error + Send + Sync>> {
        let rows = sqlx::query(
            r#"
            SELECT id, command_type, target, payload, status, results FROM commands WHERE target = $1 ORDER BY created_at DESC
            "#
        )
        .bind(target)
        .fetch_all(&self.pool)
        .await?;
        
        let mut commands = Vec::new();
        for row in rows {
            let command = CommandRecord {
                id: row.get("id"),
                command_type: match row.get::<&str, _>("command_type") {
                    "shell" => CommandType::Shell,
                    "script" => CommandType::Script,
                    "service_restart" => CommandType::ServiceRestart,
                    "file_operation" => CommandType::FileOperation,
                    _ => CommandType::Shell,
                },
                target: row.get("target"),
                payload: row.get("payload"),
                status: match row.get::<&str, _>("status") {
                    "pending" => CommandStatus::Pending,
                    "running" => CommandStatus::Running,
                    "completed" => CommandStatus::Completed,
                    "failed" => CommandStatus::Failed,
                    "cancelled" => CommandStatus::Cancelled,
                    _ => CommandStatus::Pending,
                },
                created_at: Utc::now(), // 简化处理
                completed_at: None,
                results: serde_json::from_value(row.get("results"))?,
            };
            commands.push(command);
        }
        
        Ok(commands)
    }
    
    // 文件管理实现（简化版本）
    async fn store_file_metadata(&self, metadata: &FileMetadata) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        sqlx::query(
            r#"
            INSERT INTO files (id, filename, file_size, checksum, status)
            VALUES ($1, $2, $3, $4, $5)
            "#
        )
        .bind(&metadata.id)
        .bind(&metadata.filename)
        .bind(metadata.file_size)
        .bind(&metadata.checksum)
        .bind(match metadata.status {
            FileStatus::Pending => "pending",
            FileStatus::Transferring => "transferring",
            FileStatus::Completed => "completed",
            FileStatus::Failed => "failed",
        })
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn update_file_status(&self, file_id: &str, status: FileStatus) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let status_str = match status {
            FileStatus::Pending => "pending",
            FileStatus::Transferring => "transferring",
            FileStatus::Completed => "completed",
            FileStatus::Failed => "failed",
        };
        
        sqlx::query(
            "UPDATE files SET status = $1, updated_at = NOW() WHERE id = $2"
        )
        .bind(status_str)
        .bind(file_id)
        .execute(&self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn get_file_metadata(&self, file_id: &str) -> Result<Option<FileMetadata>, Box<dyn std::error::Error + Send + Sync>> {
        let row = sqlx::query(
            "SELECT id, filename, file_size, checksum, status FROM files WHERE id = $1"
        )
        .bind(file_id)
        .fetch_optional(&self.pool)
        .await?;
        
        if let Some(row) = row {
            let metadata = FileMetadata {
                id: row.get("id"),
                filename: row.get("filename"),
                file_size: row.get("file_size"),
                checksum: row.get("checksum"),
                status: match row.get::<&str, _>("status") {
                    "pending" => FileStatus::Pending,
                    "transferring" => FileStatus::Transferring,
                    "completed" => FileStatus::Completed,
                    "failed" => FileStatus::Failed,
                    _ => FileStatus::Pending,
                },
                created_at: Utc::now(), // 简化处理
                updated_at: Utc::now(),
            };
            Ok(Some(metadata))
        } else {
            Ok(None)
        }
    }
}

impl std::str::FromStr for NodeStatus {
    type Err = crate::common::errors::Error;
    
    fn from_str(s: &str) -> Result<Self> {
        match s.to_uppercase().as_str() {
            "UNKNOWN" => Ok(NodeStatus::Unknown),
            "ONLINE" => Ok(NodeStatus::Online),
            "OFFLINE" => Ok(NodeStatus::Offline),
            "MAINTENANCE" => Ok(NodeStatus::Maintenance),
            "ERROR" => Ok(NodeStatus::Error),
            _ => Err(crate::common::errors::Error::InvalidInput(format!("Invalid node status: {}", s))),
        }
    }
} 